#include <cstring>
#include <string>
#include <iostream>
#include <fstream>
#include <chrono>
#include <iomanip>
#include <sys/time.h>
#include <omp.h>
// 可以自行添加需要的头文件
#include <cmath>
#include <cassert>
#include <algorithm>
#include <arm_neon.h>
#include <stdint.h>
#include <stdlib.h>
#include <cstdint>

void fRead(int *a, int *b, int *n, int *p, int input_id){
    // 数据输入函数
    std::string str1 = "/nttdata/";
    std::string str2 = std::to_string(input_id);
    std::string strin = str1 + str2 + ".in";
    char data_path[strin.size() + 1];
    std::copy(strin.begin(), strin.end(), data_path);
    data_path[strin.size()] = '\0';
    std::ifstream fin;
    fin.open(data_path, std::ios::in);
    fin>>*n>>*p;
    for (int i = 0; i < *n; i++){
        fin>>a[i];
    }
    for (int i = 0; i < *n; i++){   
        fin>>b[i];
    }
}

void fCheck(int *ab, int n, int input_id){
    // 判断多项式乘法结果是否正确
    std::string str1 = "/nttdata/";
    std::string str2 = std::to_string(input_id);
    std::string strout = str1 + str2 + ".out";
    char data_path[strout.size() + 1];
    std::copy(strout.begin(), strout.end(), data_path);
    data_path[strout.size()] = '\0';
    std::ifstream fin;
    fin.open(data_path, std::ios::in);
    for (int i = 0; i < n * 2 - 1; i++){
        int x;
        fin>>x;
        if(x != ab[i]){
            std::cout<<"多项式乘法结果错误"<<std::endl;
            return;
        }
    }
    std::cout<<"多项式乘法结果正确"<<std::endl;
    return;
}

void fWrite(int *ab, int n, int input_id){
    // 数据输出函数, 可以用来输出最终结果, 也可用于调试时输出中间数组
    std::string str1 = "files/";
    std::string str2 = std::to_string(input_id);
    std::string strout = str1 + str2 + ".out";
    char output_path[strout.size() + 1];
    std::copy(strout.begin(), strout.end(), output_path);
    output_path[strout.size()] = '\0';
    std::ofstream fout;
    fout.open(output_path, std::ios::out);
    for (int i = 0; i < n * 2 - 1; i++){
        fout<<ab[i]<<'\n';
    }
}

//-------------------辅助设计-----------------------------
int qpow(int a,int b,int p){//快速幂
  int res = 1;
  while (b) {
    if (b & 1) res = 1LL * res * a % p;
    a = 1LL * a * a % p;
    b >>= 1;
    }
  return res;
}

static inline int32x4_t neon_mod_mul(int32x4_t a, int32x4_t b, int32x4_t p) {// NEON向量化模乘辅助函数
    // 手动处理每个元素的模乘运算
    int32_t a_array[4], b_array[4], p_array[4], result[4];

    // 将向量存储到数组中
    vst1q_s32(a_array, a);
    vst1q_s32(b_array, b);
    vst1q_s32(p_array, p);

    // 对每个元素执行模乘运算
    for (int i = 0; i < 4; i++) {
        result[i] = (1LL * a_array[i] * b_array[i]) % p_array[i];
    }

    // 将结果加载回向量
    return vld1q_s32(result);
}

// 模运算
static inline int32x4_t neon_mod(int32x4_t a, int32x4_t p) {
    int32_t a_array[4], p_array[4], result[4];

    vst1q_s32(a_array, a);
    vst1q_s32(p_array, p);

    for (int i = 0; i < 4; i++) {
        result[i] = a_array[i] % p_array[i];
        if (result[i] < 0) result[i] += p_array[i];
    }

    return vld1q_s32(result);
}

struct MontgomeryParams {
    uint32_t mod;
    uint32_t n_prime;
    uint32_t r_squared;
};

static inline MontgomeryParams init_montgomery_params(uint32_t modulus) {
    MontgomeryParams params;
    params.mod = modulus;

    uint32_t x = 1;
    for (int i = 0; i < 5; i++) x = x * (2 - modulus * x);
    params.n_prime = -x;

    uint64_t r_squared = (1ULL << 32) % modulus;
    r_squared = (r_squared * r_squared) % modulus;
    params.r_squared = (uint32_t)r_squared;
    return params;
}

// ------------------- Montgomery运算 ----------------------
static inline uint32_t montgomery_reduce_scalar(uint64_t T, const MontgomeryParams* params) {
    uint32_t m = (uint32_t)T * params->n_prime;
    uint64_t t = (T + (uint64_t)m * params->mod) >> 32;
    return (t >= params->mod) ? (t - params->mod) : t;
}

static inline uint32_t to_montgomery_scalar(uint32_t a, const MontgomeryParams* params) {
    return montgomery_reduce_scalar((uint64_t)a * params->r_squared, params);
}

static inline uint32_t from_montgomery_scalar(uint32_t a, const MontgomeryParams* params) {
    return montgomery_reduce_scalar((uint64_t)a, params);
}

static inline uint32x4_t montgomery_reduce_neon(uint64x2_t T_low, uint64x2_t T_high, const MontgomeryParams* params) {
    uint32x4_t n_prime_vec = vdupq_n_u32(params->n_prime);
    uint32x4_t mod_vec = vdupq_n_u32(params->mod);
    uint32x4_t T_low32 = vcombine_u32(vmovn_u64(T_low), vmovn_u64(T_high));
    uint32x4_t m = vmulq_u32(T_low32, n_prime_vec);

    uint64x2_t mul_low = vmull_u32(vget_low_u32(m), vget_low_u32(mod_vec));
    uint64x2_t mul_high = vmull_u32(vget_high_u32(m), vget_high_u32(mod_vec));
    uint64x2_t sum_low = vaddq_u64(T_low, mul_low);
    uint64x2_t sum_high = vaddq_u64(T_high, mul_high);

    uint32x4_t t = vcombine_u32(vshrn_n_u64(sum_low, 32), vshrn_n_u64(sum_high, 32));
    uint32x4_t cmp = vcgeq_u32(t, mod_vec);
    uint32x4_t sub = vsubq_u32(t, mod_vec);
    return vbslq_u32(cmp, sub, t);
}

static inline uint32x4_t montgomery_multiply_neon(uint32x4_t a, uint32x4_t b, const MontgomeryParams* params) {
    uint64x2_t low = vmull_u32(vget_low_u32(a), vget_low_u32(b));
    uint64x2_t high = vmull_u32(vget_high_u32(a), vget_high_u32(b));
    return montgomery_reduce_neon(low, high, params);
}

static void prepare_w_table(uint32_t* W, int len, uint32_t wn, uint32_t p, const MontgomeryParams* params) {
    W[0] = to_montgomery_scalar(1, params); // 1的Montgomery表示
    for (int i = 1; i < len; i++) {
        // Montgomery 域乘法
        W[i] = montgomery_reduce_scalar((uint64_t)W[i-1] * wn, params);
    }
}

uint32_t montgomery_pow(uint32_t a, int b, const MontgomeryParams* params) {
    uint32_t res = to_montgomery_scalar(1, params); // Montgomery域的1
    while (b) {
        if (b & 1) res = montgomery_reduce_scalar((uint64_t)res * a, params);
        a = montgomery_reduce_scalar((uint64_t)a * a, params);
        b >>= 1;
    }
    return res;
}

//-------------------算法部分-----------------------------
void poly_multiply(int *a, int *b, int *ab, int n, int p){
    for(int i = 0; i < n; ++i){
        for(int j = 0; j < n; ++j){
            ab[i+j]=(1LL * a[i] * b[j] % p + ab[i+j]) % p;
        }
    }
}

void poly_multiply_neon(int *a, int *b, int *ab, int n, int p) {//朴素乘法的SIMD优化
    // 初始化结果数组
    for (int i = 0; i < 2 * n - 1; ++i) {
        ab[i] = 0;
    }

    for (int i = 0; i < n; ++i) {
        int a_val = a[i];
        int32x2_t a_low2 = vdup_n_s32(a_val);  // 复制两个 a[i]

        for (int j = 0; j + 3 < n; j += 4) {
            // 加载 b[j], b[j+1], b[j+2], b[j+3]
            int32x4_t b_vec = vld1q_s32(&b[j]);

            // 拆成低2位和高2位
            int32x2_t b_low2 = vget_low_s32(b_vec);
            int32x2_t b_high2 = vget_high_s32(b_vec);

            // 乘法：a[i] * b[j~j+3] -> int64x2_t
            int64x2_t mul_low2 = vmull_s32(a_low2, b_low2);   // a[i] * b[j], b[j+1]
            int64x2_t mul_high2 = vmull_s32(a_low2, b_high2); // a[i] * b[j+2], b[j+3]

            // 加上 ab[i+j ~ i+j+3]
            int64_t tmp[4];
            tmp[0] = mul_low2[0] + ab[i + j + 0];
            tmp[1] = mul_low2[1] + ab[i + j + 1];
            tmp[2] = mul_high2[0] + ab[i + j + 2];
            tmp[3] = mul_high2[1] + ab[i + j + 3];

            // 取模写回
            for (int k = 0; k < 4; ++k) {
                tmp[k] %= p;
                if (tmp[k] < 0) tmp[k] += p;
                ab[i + j + k] = (int)tmp[k];
            }
        }

        // 处理剩下不足4个的部分
        for (int j = (n & ~3); j < n; ++j) {
            ab[i + j] = (1LL * a[i] * b[j] % p + ab[i + j]) % p;
        }
    }
}

const int g=3;
// NTT递归实现
void NTT_recursion(int *a,int len,int op,int g,int p){//NTT递归模块
    if (len == 1) return;
    int half = len / 2;
    int A1[half], A2[half];

    for (int i = 0; i < half; ++i) {
        A1[i] = a[i * 2];
        A2[i] = a[i * 2 + 1];
    }

    NTT_recursion(A1, half, op, g, p);
    NTT_recursion(A2, half, op, g, p);

    int wn = qpow((op == 1 ? g : qpow(g, p - 2, p)), (p - 1) / len, p);
    int w = 1;

    for (int i = 0; i < half; ++i) {
        int t = 1LL * A2[i] * w % p;
        a[i] = (A1[i] + t) % p;
        a[i + half] = (A1[i] - t + p) % p;
        w = 1LL * w * wn % p;
    }
}

void NTT_recursion_multiply(int *a, int *b, int *ab, int n, int p){//NTT递归
    int N = 2 * n - 1;
    int len = 1;
    int ni;
    while (len < N) len <<= 1; // 求len,且保证为2的整数次幂
    ni=qpow(len,p-2,p);
    int A[len], B[len];
    for (int i = 0; i < len; ++i) {
        A[i] = (i < n) ? a[i] : 0;
        B[i] = (i < n) ? b[i] : 0;
    }
    NTT_recursion(A, len, 1, g, p);
    NTT_recursion(B, len, 1, g, p);
    for (int i = 0; i < len; ++i) {
        A[i] = (1LL * A[i] * B[i]) % p;
    }
    NTT_recursion(A, len, -1, g, p);
    for (int i = 0; i < N; ++i) {
        ab[i] = (1LL * A[i] * ni) % p;
    }
}

// 使用Montgomery算法neon优化的递归NTT
void NTT_recursion_mont(uint32_t *a, int len, int op, uint32_t g, uint32_t p, const MontgomeryParams *params) {
    if (len == 1) return;
    int half = len / 2;
    uint32_t *A1 = (uint32_t*)alloca(sizeof(uint32_t) * half);
    uint32_t *A2 = (uint32_t*)alloca(sizeof(uint32_t) * half);

    for (int i = 0; i < half; ++i) {
        A1[i] = a[i * 2];
        A2[i] = a[i * 2 + 1];
    }

    NTT_recursion_mont(A1, half, op, g, p, params);
    NTT_recursion_mont(A2, half, op, g, p, params);

    uint32_t root = (op == 1 ? g : qpow(g, p - 2, p));
    uint32_t root_mont = to_montgomery_scalar(root, params);
    uint32_t wn = montgomery_pow(root_mont, (p - 1) / len, params);
    uint32_t w = to_montgomery_scalar(1, params);

    for (int i = 0; i < half; ++i) {
        // t = A2[i] * w (Montgomery乘法)
        uint32_t t = montgomery_reduce_scalar((uint64_t)A2[i] * w, params);

        // a[i] = (A1[i] + t) % p;
        uint32_t add = A1[i] + t;
        if (add >= p) add -= p;
        a[i] = add;

        // a[i+half] = (A1[i] - t + p) % p;
        uint32_t sub = A1[i] + p - t;
        if (sub >= p) sub -= p;
        a[i + half] = sub;

        // w = w * wn (Montgomery)
        w = montgomery_reduce_scalar((uint64_t)w * wn, params);
    }
}

void NTT_recursion_multiply_SIMD(int *a, int *b, int *ab, int n, int p, int g=3) {
    const int N = 2 * n - 1;
    int len = 1;
    while (len < N) len <<= 1;
    int ni = qpow(len, p - 2, p);

    MontgomeryParams params = init_montgomery_params(p);

    // 转Montgomery域
    uint32_t *A = (uint32_t*)aligned_alloc(16, len * sizeof(uint32_t));
    uint32_t *B = (uint32_t*)aligned_alloc(16, len * sizeof(uint32_t));
    for (int i = 0; i < len; ++i) {
        A[i] = (i < n) ? to_montgomery_scalar(a[i], &params) : 0;
        B[i] = (i < n) ? to_montgomery_scalar(b[i], &params) : 0;
    }

    NTT_recursion_mont(A, len, 1, g, p, &params);
    NTT_recursion_mont(B, len, 1, g, p, &params);

    // 点值相乘Montgomery
    for (int i = 0; i < len; ++i) {
        A[i] = montgomery_reduce_scalar((uint64_t)A[i] * B[i], &params);
    }

    NTT_recursion_mont(A, len, -1, g, p, &params);

    for (int i = 0; i < N; ++i) {
        ab[i] = (uint64_t)from_montgomery_scalar(A[i], &params) * ni % p;
    }

    free(A);
    free(B);
}

// NTT迭代实现
void NTT_Iteration(int *a, int len, int op, int g, int p) {//NTT迭代模块
    // 位逆序置换
    for (int i = 1, j = 0; i < len - 1; i++) {
        for (int k = len >> 1; (j ^= k) < k; k >>= 1);
        if (i < j) {
            int tmp = a[i];
            a[i] = a[j];
            a[j] = tmp;
        }
    }

    // 蝴蝶操作
    for (int h = 2; h <= len; h <<= 1) {
        // 计算原根
        int wn = qpow((op == 1 ? g : qpow(g, p - 2, p)), (p - 1) / h, p);

        // 对每个长度为h的段进行变换
        for (int j = 0; j < len; j += h) {
            int w = 1;
            for (int k = j; k < j + h/2; k++) {
                int t = 1LL * a[k + h/2] * w % p;
                a[k + h/2] = (a[k] - t + p) % p;
                a[k] = (a[k] + t) % p;
                w = 1LL * w * wn % p;
            }
        }
    }
}

void NTT_Iteration_multiply(int *a, int *b, int *ab, int n, int p) {//NTT迭代
    int N = 2 * n - 1;
    int len = 1;
    while (len < N) len <<= 1;  // 确保len为2的幂次

    // 计算len的逆元
    int ni = qpow(len, p-2, p);

    // 创建并初始化临时数组
    int *A = new int[len]();
    int *B = new int[len]();

    // 复制输入数据
    for (int i = 0; i < len; i++) {
        A[i] = (i < n) ? a[i] : 0;
        B[i] = (i < n) ? b[i] : 0;
    }

    // NTT变换
    NTT_Iteration(A, len, 1, g, p);
    NTT_Iteration(B, len, 1, g, p);

    // 点值乘法
    for (int i = 0; i < len; i++) {
        A[i] = (1LL * A[i] * B[i]) % p;
    }

    // 逆变换
    NTT_Iteration(A, len, -1, g, p);

    // 乘上ni并保存结果
    for (int i = 0; i < N; i++) {
        ab[i] = (1LL * A[i] * ni) % p;
    }

    delete[] A;
    delete[] B;
}

// 使用Montgomery算法neon优化的NTT迭代实现
void NTT_Iteration_SIMD(uint32_t* a, int len, int op, uint32_t g, uint32_t p, const MontgomeryParams* params) {
    // 位逆序置换
    for (int i = 1, j = 0; i < len-1; i++) {
        for (int k = len >> 1; (j ^= k) < k; k >>= 1);
        if (i < j) { uint32_t tmp = a[i]; a[i] = a[j]; a[j] = tmp; }
    }

    uint32_t* W = (uint32_t*)aligned_alloc(16, len * sizeof(uint32_t));
    uint32x4_t mod_vec = vdupq_n_u32(p);

    for (int h = 2; h <= len; h <<= 1) {
        uint32_t root = (op == 1 ? g : qpow(g, p-2, p));
        uint32_t root_mont = to_montgomery_scalar(root, params); // 先转域
        uint32_t wn = montgomery_pow(root_mont, (p-1)/h, params); // 全在 Montgomery 域
        prepare_w_table(W, h/2, wn, p, params);

        for (int j = 0; j < len; j += h) {
            int k = 0;
            for (; k + 4 <= h/2; k += 4) {
                uint32x4_t vA = vld1q_u32(a + j + k);
                uint32x4_t vB = vld1q_u32(a + j + k + h/2);
                uint32x4_t vW = vld1q_u32(W + k);
                uint32x4_t vT = montgomery_multiply_neon(vB, vW, params);

                uint32x4_t vSum = vaddq_u32(vA, vT);
                uint32x4_t vDiff = vsubq_u32(vA, vT);

                vSum = vbslq_u32(vcgeq_u32(vSum, mod_vec), vsubq_u32(vSum, mod_vec), vSum);
                vDiff = vbslq_u32(vcgtq_u32(vT, vA), vaddq_u32(vDiff, mod_vec), vDiff);

                vst1q_u32(a + j + k, vSum);
                vst1q_u32(a + j + k + h/2, vDiff);
            }
            for (; k < h/2; ++k) {
                uint32_t t = montgomery_reduce_scalar((uint64_t)a[j+k+h/2] * W[k], params);
                uint32_t u = a[j+k];
                uint32_t add = u + t;
                if (add >= p) add -= p;
                a[j+k] = add;
                uint32_t sub = u + p - t;
                if (sub >= p) sub -= p;
                a[j+k+h/2] = sub;
            }
        }
    }
    free(W);
}

void NTT_Iteration_multiply_SIMD(int* a, int* b, int* ab, int n, int p) {
    const int N = 2*n-1;
    int len = 1;
    while (len < N) len <<= 1;
    int ni = qpow(len, p-2, p);

    uint32_t* A = (uint32_t*)aligned_alloc(16, len * sizeof(uint32_t));
    uint32_t* B = (uint32_t*)aligned_alloc(16, len * sizeof(uint32_t));

    MontgomeryParams params = init_montgomery_params(p);

    // 输入数据 to_montgomery_scalar
    for (int i = 0; i < len; ++i) {
        A[i] = (i < n) ? to_montgomery_scalar(a[i], &params) : 0;
        B[i] = (i < n) ? to_montgomery_scalar(b[i], &params) : 0;
    }

    // 正向NTT
    NTT_Iteration_SIMD(A, len, 1, 3, p, &params);
    NTT_Iteration_SIMD(B, len, 1, 3, p, &params);

    // 点值乘
    for (int i = 0; i < len; ++i) {
        A[i] = montgomery_reduce_scalar((uint64_t)A[i] * B[i], &params);
    }

    // 逆向NTT
    NTT_Iteration_SIMD(A, len, -1, 3, p, &params);

    // 输出 from_montgomery_scalar
    for (int i = 0; i < N; ++i) {
        ab[i] = (uint64_t)from_montgomery_scalar(A[i], &params) * ni % p;
    }

    free(A);
    free(B);
}

int a[300000], b[300000], ab[300000];
int main(int argc, char *argv[])
{

    // 保证输入的所有模数的原根均为 3, 且模数都能表示为 a \times 4 ^ k + 1 的形式
    // 输入模数分别为 7340033 104857601 469762049 263882790666241
    // 第四个模数超过了整型表示范围, 如果实现此模数意义下的多项式乘法需要修改框架
    // 对第四个模数的输入数据不做必要要求, 如果要自行探索大模数 NTT, 请在完成前三个模数的基础代码及优化后实现大模数 NTT
    // 输入文件共五个, 第一个输入文件 n = 4, 其余四个文件分别对应四个模数, n = 131072
    // 在实现快速数论变化前, 后四个测试样例运行时间较久, 推荐调试正确性时只使用输入文件 1
    int test_begin = 0;
    int test_end = 3;
    for(int i = test_begin; i <= test_end; ++i){
        long double ans = 0;
        int n_, p_;
        fRead(a, b, &n_, &p_, i);
        memset(ab,0,sizeof(ab));
        auto Start = std::chrono::high_resolution_clock::now();
        // TODO : 将 poly_multiply 函数替换成你写的 ntt
        // poly_multiply(a, b, ab, n_, p_);
        // poly_multiply_neon(a, b, ab, n_, p_);
        // NTT_recursion_multiply(a, b, ab, n_, p_);
        // NTT_recursion_multiply_SIMD(a, b, ab, n_, p_);
        // NTT_Iteration_multiply(a, b, ab, n_, p_);
        NTT_Iteration_multiply_SIMD(a, b, ab, n_, p_);
        auto End = std::chrono::high_resolution_clock::now();
        std::chrono::duration<double,std::ratio<1,1000>>elapsed = End - Start;
        ans += elapsed.count();
        fCheck(ab, n_, i);
        std::cout<<"average latency for n = "<<n_<<" p = "<<p_<<" : "<<ans<<" (us) "<<std::endl;
        // 可以使用 fWrite 函数将 ab 的输出结果打印到 files 文件夹下
        // 禁止使用 cout 一次性输出大量文件内容
        fWrite(ab, n_, i);
    }
    return 0;
}